import turtle
import pandas as pd
import numpy as np
import random
import time
import os


# 动物基类
class Animal:
    def __init__(self, name, color, shape, speed):
        self.name = name
        self.color = color  # 角色颜色，用于可视化
        self.shape = shape
        self.speed = speed
        self.turtle = turtle.Turtle()
        self.turtle.shape(shape)
        self.turtle.color(color)
        self.turtle.speed(0)
        self.turtle.penup()
        self.trajectory = []  # 轨迹数据：(x, y, 时间戳)
        self.active = True  # 活动状态标记

    def move_forward(self, distance):
        self.turtle.forward(distance)
        x, y = self.turtle.position()
        self.trajectory.append((x, y, time.time()))

    def turn(self, angle):
        self.turtle.right(angle)

    def start_trace(self):
        self.turtle.pendown()

    def stop_trace(self):
        self.turtle.penup()

    def get_position(self):
        return self.turtle.position()

    def get_trajectory_dataframe(self):
        """将轨迹数据转换为DataFrame"""
        df = pd.DataFrame(self.trajectory, columns=['x', 'y', 'timestamp'])
        df['animal'] = self.name
        df['color'] = self.color  # 颜色信息写入DataFrame
        return df

    def is_collision(self, other_animal, distance_threshold=20):
        """检测与其他动物是否碰撞"""
        x1, y1 = self.get_position()
        x2, y2 = other_animal.get_position()
        distance = np.sqrt((x2 - x1) ** 2 + (y2 - y1) ** 2)
        return distance < distance_threshold


# 人类类
class Human(Animal):
    def __init__(self, name="人类", color="purple", speed=2):
        super().__init__(name, color, "circle", speed)
        self.score = 0  # 生存计分

    def move_strategy(self, crocodiles=None):
        """人类移动策略：优先远离鳄鱼"""
        if crocodiles and random.random() < 0.7:
            hx, hy = self.get_position()
            dx_sum, dy_sum = 0, 0
            for croc in crocodiles:
                cx, cy = croc.get_position()
                dx_sum += hx - cx
                dy_sum += hy - cy

            distance = np.sqrt(dx_sum ** 2 + dy_sum ** 2)
            if distance < 150:
                angle = np.degrees(np.arctan2(dy_sum, dx_sum))
                self.turtle.setheading(angle)
                self.move_forward(self.speed * 1.5)
                return

        # 随机移动逻辑
        if random.random() < 0.2:
            self.turn(random.randint(-90, 90))
        self.move_forward(self.speed)
        self.score += 1


# 鳄鱼类
class Crocodile(Animal):
    def __init__(self, name="鳄鱼", color="green", speed=1.5):
        super().__init__(name, color, "triangle", speed)
        self.hunger = 0  # 饥饿值

    def move_strategy(self, human=None):
        """鳄鱼移动策略：优先追踪人类"""
        if human and random.random() < 0.8:
            hx, hy = human.get_position()
            cx, cy = self.get_position()
            dx = hx - cx
            dy = hy - cy

            angle = np.degrees(np.arctan2(dy, dx))
            self.turtle.setheading(angle)
            distance = np.sqrt(dx ** 2 + dy ** 2)

            if distance < 100:
                self.move_forward(self.speed * 1.2)
                self.hunger -= 0.5
            else:
                self.move_forward(self.speed)
                self.hunger += 0.1
            return

        # 随机移动逻辑
        if random.random() < 0.3:
            self.turn(random.randint(-60, 60))
        self.move_forward(self.speed)
        self.hunger += 0.2
        self.hunger = min(self.hunger, 100)

    def is_hungry(self):
        """判断鳄鱼是否饥饿"""
        return self.hunger > 70


# 环境类
class Environment:
    def __init__(self, width=800, height=600, background_color="lightblue"):
        self.width = width
        self.height = height
        self.screen = turtle.Screen()
        self.screen.title("人与鳄鱼追踪模拟")
        self.screen.setup(width, height)
        self.screen.bgcolor(background_color)
        self.screen.tracer(0)

        self.human = Human()
        self.crocodiles = []
        self._create_crocodiles()  # 初始化3只鳄鱼
        self.obstacles = []
        self.animals = [self.human] + self.crocodiles

        # 分数显示
        self.score_display = turtle.Turtle()
        self.score_display.speed(0)
        self.score_display.color("black")
        self.score_display.penup()
        self.score_display.hideturtle()
        self.score_display.goto(-self.width / 2 + 50, self.height / 2 - 30)

        # 创建障碍物
        self._create_obstacles(10)

    def _create_crocodiles(self):
        """初始化3只鳄鱼，等边三角形分布"""
        center_x, center_y = 0, 0  # 人类初始在中心
        radius = 150  # 等边三角形外接圆半径
        angles = [0, 120, 240]  # 三个顶点角度
        colors = ["green", "blue", "orange"]  # 每只鳄鱼颜色
        speeds = [2, 2, 1.5]  # 前两只速度2，第三只1.5

        for i in range(3):
            angle_rad = np.radians(angles[i])
            x = center_x + radius * np.cos(angle_rad)
            y = center_y + radius * np.sin(angle_rad)

            # 初始化鳄鱼
            croc = Crocodile(
                name=f"鳄鱼{i + 1}",
                color=colors[i],
                speed=speeds[i]
            )
            croc.turtle.goto(x, y)
            self.crocodiles.append(croc)

    def _create_obstacles(self, count):
        """随机生成障碍物"""
        for _ in range(count):
            obstacle = turtle.Turtle()
            obstacle.shape("square")
            obstacle.color("brown")
            obstacle.penup()

            size = random.randint(10, 30)
            obstacle.shapesize(stretch_len=size / 20, stretch_wid=size / 20)

            x = random.randint(-self.width // 2 + 50, self.width // 2 - 50)
            y = random.randint(-self.height // 2 + 50, self.height // 2 - 50)
            obstacle.goto(x, y)

            self.obstacles.append((obstacle, size))

    def check_boundaries(self, animal):
        """边界检测：防止动物超出屏幕范围"""
        x, y = animal.get_position()

        # X轴边界
        if x > self.width / 2 - 20:
            animal.turtle.setx(self.width / 2 - 20)
            animal.turn(180)
        elif x < -self.width / 2 + 20:
            animal.turtle.setx(-self.width / 2 + 20)
            animal.turn(180)

        # Y轴边界
        if y > self.height / 2 - 20:
            animal.turtle.sety(self.height / 2 - 20)
            animal.turn(180)
        elif y < -self.height / 2 + 20:
            animal.turtle.sety(-self.height / 2 + 20)
            animal.turn(180)

    def check_obstacle_collision(self, animal):
        """障碍物碰撞检测：碰到障碍物则转向"""
        x, y = animal.get_position()
        for obstacle, size in self.obstacles:
            ox, oy = obstacle.position()
            distance = np.sqrt((x - ox) ** 2 + (y - oy) ** 2)

            if distance < (size + 10):
                animal.turn(random.randint(90, 270))
                return True
        return False

    def update_score_display(self):
        """更新分数显示"""
        self.score_display.clear()
        score_text = (
            f"人类存活时间: {self.human.score // 10}s | "
            f"鳄鱼1饥饿值: {int(self.crocodiles[0].hunger)} | "
            f"鳄鱼2饥饿值: {int(self.crocodiles[1].hunger)} | "
            f"鳄鱼3饥饿值: {int(self.crocodiles[2].hunger)}"
        )
        self.score_display.write(score_text, font=("Arial", 12, "normal"))

    def run_simulation(self, steps=1000, record_interval=10):
        """运行模拟：任一鳄鱼追上人类即结束"""
        self.human.turtle.goto(0, 0)  # 人类初始在中心

        # 开始绘制轨迹
        self.human.start_trace()
        for croc in self.crocodiles:
            croc.start_trace()

        trajectory_data = []  # 保存轨迹数据

        for step in range(steps):
            if not self.human.active:
                print("人类被鳄鱼抓住了！")
                break

            # 1. 执行移动策略
            self.human.move_strategy(self.crocodiles)
            for croc in self.crocodiles:
                croc.move_strategy(self.human)

            # 2. 边界检测
            self.check_boundaries(self.human)
            for croc in self.crocodiles:
                self.check_boundaries(croc)

            # 3. 障碍物碰撞检测
            self.check_obstacle_collision(self.human)
            for croc in self.crocodiles:
                self.check_obstacle_collision(croc)

            # 4. 捕获检测（任一鳄鱼追上即结束）
            for croc in self.crocodiles:
                if croc.is_collision(self.human) and croc.is_hungry():
                    self.human.active = False
                    self.human.turtle.hideturtle()
                    self.human.stop_trace()

                    # 所有鳄鱼停止绘制轨迹
                    for c in self.crocodiles:
                        c.stop_trace()
                    break  # 触发结束

            # 5. 记录轨迹（按间隔采样）
            if step % record_interval == 0:
                human_data = self.human.get_trajectory_dataframe().iloc[-1:]
                croc_datas = [
                    croc.get_trajectory_dataframe().iloc[-1:]
                    for croc in self.crocodiles
                ]
                trajectory_data.append(pd.concat([human_data] + croc_datas))

            # 6. 更新界面
            self.update_score_display()
            self.screen.update()
            time.sleep(0.01)

        # 拼接轨迹数据
        return pd.concat(trajectory_data) if trajectory_data else pd.DataFrame()

    def save_trajectory_to_csv(self, data, filename="animal_trajectory.csv"):
        """保存轨迹数据到CSV文件"""
        if not data.empty:
            data.to_csv(filename, index=False)
            print(f"轨迹数据已保存到 {filename}")
        else:
            print("无轨迹数据可保存")

    def close(self):
        """关闭模拟窗口"""
        self.screen.bye()


# 轨迹分析与可视化
def analyze_trajectory_data(csv_file="animal_trajectory.csv"):
    """分析并可视化轨迹数据"""
    if not os.path.exists(csv_file):
        print(f"文件不存在: {csv_file}")
        return

    df = pd.read_csv(csv_file)
    print("\n轨迹分析结果:")
    print(f"总记录数: {len(df)}")
    print(f"人类 X 方向移动标准差: {df[df['animal'] == '人类']['x'].std():.2f}")

    for i in range(3):
        croc_name = f"鳄鱼{i + 1}"
        print(f"{croc_name} X 方向移动标准差: {df[df['animal'] == croc_name]['x'].std():.2f}")

    # Plotly 可视化（轨迹颜色与角色颜色一致）
    try:
        import plotly.express as px

        fig = px.scatter(
            df,
            x="x",
            y="y",
            color="color",  # 用角色颜色作为轨迹颜色
            animation_frame="timestamp",
            title="人与鳄鱼动态轨迹",
            labels={"x": "X坐标", "y": "Y坐标"},
            range_x=[-450, 450],
            range_y=[-350, 350],
            color_discrete_map={  # 确保颜色映射一致
                "purple": "purple",
                "green": "green",
                "blue": "blue",
                "orange": "orange"
            }
        )

        # 添加轨迹连线
        for animal in df['animal'].unique():
            animal_df = df[df['animal'] == animal]
            fig.add_trace(
                px.line(animal_df, x="x", y="y",
                        color_discrete_sequence=[animal_df['color'].iloc[0]]).data[0]
            )

        fig.show()

    except ImportError:
        print("缺少 Plotly，请运行 'pip install plotly' 安装可视化依赖")


# 主程序
def main():
    env = Environment()
    try:
        print("开始模拟...")
        trajectory_data = env.run_simulation(steps=500)
        env.save_trajectory_to_csv(trajectory_data)
        analyze_trajectory_data()
    except Exception as e:
        print(f"模拟出错: {e}")
    finally:
        print("点击窗口关闭模拟...")
        env.screen.exitonclick()


if __name__ == "__main__":
    main()